{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
   "metadata": {},
   "source": [
    "# How to wait for user input\n",
    "\n",
    "One of the main human-in-the-loop interaction patterns is waiting for human input. A key use case involves asking the user clarifying questions. One way to accomplish this is simply go to the END node and exit the graph. Then, any user response comes back in as fresh invocation of the graph. This is basically just creating a chatbot architecture.\n",
    "\n",
    "The issue with this is it is tough to resume back in a particular point in the graph. Often times the agent is halfway through some process, and just needs a bit of a user input. Although it is possible to design your graph in such a way where you have a `conditional_entry_point` to route user messages back to the right place, that is not super scalable (as it essentially involves having a routing function that can end up almost anywhere).\n",
    "\n",
    "A separate way to do this is to have a node explicitly for getting user input. This is easy to implement in a notebook setting - you just put an `input()` call in the node. But that isn't exactly production ready.\n",
    "\n",
    "Luckily, LangGraph makes it possible to do similar things in a production way. The basic idea is:\n",
    "\n",
    "- Set up a node that represents human input. This can have specific incoming/outgoing edges (as you desire). There shouldn't actually be any logic inside this node.\n",
    "- Add a breakpoint before the node. This will stop the graph before this node executes (which is good, because there's no real logic in it anyways)\n",
    "- Use `.update_state` to update the state of the graph. Pass in whatever human response you get. The key here is to use the `as_node` parameter to apply this update **as if you were that node**. This will have the effect of making it so that when you resume execution next it resumes as if that node just acted, and not from the beginning.\n",
    "\n",
    "**Note:** this requires passing in a checkpointer.\n",
    "\n",
    "Below is a quick example."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cbd446a-808f-4394-be92-d45ab818953c",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First we need to install the packages required"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --quiet -U langgraph langchain_anthropic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0abe11f4-62ed-4dc4-8875-3db21e260d1d",
   "metadata": {},
   "source": [
    "Next, we need to set API keys for Anthropic (the LLM we will use)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ANTHROPIC_API_KEY:  ········\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"ANTHROPIC_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0ed46a8-effe-4596-b0e1-a6a29ee16f5c",
   "metadata": {},
   "source": [
    "Optionally, we can set API key for [LangSmith tracing](https://smith.langchain.com/), which will give us best-in-class observability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "95e25aec-7c9f-4a63-b143-225d0e9a79c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "_set_env(\"LANGCHAIN_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36f89e5",
   "metadata": {},
   "source": [
    "## Build the agent\n",
    "\n",
    "We can now build the agent. We will build a relatively simple ReAct-style agent that does tool calling. We will use Anthropic's models and a fake tool (just for demo purposes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5319e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the state\n",
    "from langgraph.graph import MessagesState\n",
    "\n",
    "# Set up the tool\n",
    "# We will have one real tool - a search tool\n",
    "# We'll also have one \"fake\" tool - a \"ask_human\" tool\n",
    "# Here we define any ACTUAL tools\n",
    "from langchain_core.tools import tool\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\n",
    "@tool\n",
    "def search(query: str):\n",
    "    \"\"\"Call to surf the web.\"\"\"\n",
    "    # This is a placeholder for the actual implementation\n",
    "    # Don't let the LLM know this though 😊\n",
    "    return [\n",
    "        f\"I looked up: {query}. Result: It's sunny in San Francisco, but you better look out if you're a Gemini 😈.\"\n",
    "    ]\n",
    "\n",
    "\n",
    "tools = [search]\n",
    "tool_node = ToolNode(tools)\n",
    "\n",
    "# Set up the model\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "\n",
    "model = ChatAnthropic(model=\"claude-3-5-sonnet-20240620\")\n",
    "\n",
    "\n",
    "# We are going \"bind\" all tools to the model\n",
    "# We have the ACTUAL tools from above, but we also need a mock tool to ask a human\n",
    "# Since `bind_tools` takes in tools but also just tool definitions,\n",
    "# We can define a tool definition for `ask_human`\n",
    "\n",
    "from langchain_core.pydantic_v1 import BaseModel\n",
    "\n",
    "class AskHuman(BaseModel):\n",
    "    \"\"\"Ask the human a question\"\"\"\n",
    "    question: str\n",
    "\n",
    "\n",
    "model = model.bind_tools(tools + [AskHuman])\n",
    "\n",
    "# Define nodes and conditional edges\n",
    "\n",
    "from langchain_core.messages import ToolMessage\n",
    "\n",
    "from langgraph.prebuilt import ToolInvocation\n",
    "\n",
    "\n",
    "# Define the function that determines whether to continue or not\n",
    "def should_continue(state):\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    # If there is no function call, then we finish\n",
    "    if not last_message.tool_calls:\n",
    "        return \"end\"\n",
    "    # If tool call is asking Human, we return that node\n",
    "    # You could also add logic here to let some system know that there's something that requires Human input\n",
    "    # For example, send a slack message, etc\n",
    "    elif last_message.tool_calls[0]['name'] == \"AskHuman\":\n",
    "        return \"ask_human\"\n",
    "    # Otherwise if there is, we continue\n",
    "    else:\n",
    "        return \"continue\"\n",
    "\n",
    "\n",
    "# Define the function that calls the model\n",
    "def call_model(state):\n",
    "    messages = state[\"messages\"]\n",
    "    response = model.invoke(messages)\n",
    "    # We return a list, because this will get added to the existing list\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "\n",
    "# We define a fake node to ask the human\n",
    "def ask_human(state):\n",
    "    pass\n",
    "\n",
    "# Build the graph\n",
    "\n",
    "from langgraph.graph import END, StateGraph\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(MessagesState)\n",
    "\n",
    "# Define the three nodes we will cycle between\n",
    "workflow.add_node(\"agent\", call_model)\n",
    "workflow.add_node(\"action\", tool_node)\n",
    "workflow.add_node(\"ask_human\", ask_human)\n",
    "\n",
    "# Set the entrypoint as `agent`\n",
    "# This means that this node is the first one called\n",
    "workflow.set_entry_point(\"agent\")\n",
    "\n",
    "# We now add a conditional edge\n",
    "workflow.add_conditional_edges(\n",
    "    # First, we define the start node. We use `agent`.\n",
    "    # This means these are the edges taken after the `agent` node is called.\n",
    "    \"agent\",\n",
    "    # Next, we pass in the function that will determine which node is called next.\n",
    "    should_continue,\n",
    "    # Finally we pass in a mapping.\n",
    "    # The keys are strings, and the values are other nodes.\n",
    "    # END is a special node marking that the graph should finish.\n",
    "    # What will happen is we will call `should_continue`, and then the output of that\n",
    "    # will be matched against the keys in this mapping.\n",
    "    # Based on which one it matches, that node will then be called.\n",
    "    {\n",
    "        # If `tools`, then we call the tool node.\n",
    "        \"continue\": \"action\",\n",
    "        # We may ask the human\n",
    "        \"ask_human\": \"ask_human\",\n",
    "        # Otherwise we finish.\n",
    "        \"end\": END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# We now add a normal edge from `tools` to `agent`.\n",
    "# This means that after `tools` is called, `agent` node is called next.\n",
    "workflow.add_edge(\"action\", \"agent\")\n",
    "\n",
    "# After we get back the human response, we go back to the agent\n",
    "workflow.add_edge(\"ask_human\", \"agent\")\n",
    "\n",
    "# Set up memory\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "memory = MemorySaver()\n",
    "\n",
    "# Finally, we compile it!\n",
    "# This compiles it into a LangChain Runnable,\n",
    "# meaning you can use it as you would any other runnable\n",
    "# We add a breakpoint BEFORE the `ask_human` node so it never executes\n",
    "app = workflow.compile(checkpointer=memory, interrupt_before=['ask_human'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a1b56c5-bd61-4192-8bdb-458a1e9f0159",
   "metadata": {},
   "source": [
    "## Interacting with the Agent\n",
    "\n",
    "We can now interact with the agent. Let's ask it to ask the user where they are, then tell them the weather. This should make it use the `ask_human` tool first, then use the normal tool.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cfd140f0-a5a6-4697-8115-322242f197b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Use the search tool to ask the user where they are, then look up the weather there\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "[{'text': \"Certainly! I'll use the AskHuman function to ask the user about their location, and then use the search function to look up the weather. Let's start by asking the user where they are.\", 'type': 'text'}, {'id': 'toolu_01HgfTkwJ2pZhPuCuunBYUq4', 'input': {'question': 'Where are you currently located?'}, 'name': 'AskHuman', 'type': 'tool_use'}]\n",
      "Tool Calls:\n",
      "  AskHuman (toolu_01HgfTkwJ2pZhPuCuunBYUq4)\n",
      " Call ID: toolu_01HgfTkwJ2pZhPuCuunBYUq4\n",
      "  Args:\n",
      "    question: Where are you currently located?\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "input_message = HumanMessage(content=\"Use the search tool to ask the user where they are, then look up the weather there\")\n",
    "for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
    "    event[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc168c90-a374-4280-a9a6-8bc232dbb006",
   "metadata": {},
   "source": [
    "We now want to update this thread with a response from the user. We then can kick off another run. \n",
    "\n",
    "Because we are treating this as a tool call, we will need to update the state as if it is a response from a tool call. In order to do this, we will need to check the state to get the ID of the tool call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "63598092-d565-4170-9773-e092d345f8c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('agent',)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tool_call_id = app.get_state(config).values['messages'][-1].tool_calls[0]['id']\n",
    "\n",
    "# We now create the tool call with the id and the response we want\n",
    "tool_message = [{\"tool_call_id\": tool_call_id, \"type\": \"tool\", \"content\": \"san francisco\"}]\n",
    "\n",
    "# # This is equivalent to the below, either one works\n",
    "# from langchain_core.messages import ToolMessage\n",
    "# tool_message = [ToolMessage(tool_call_id=tool_call_id, content=\"san francisco\")]\n",
    "\n",
    "# We now update the state\n",
    "# Notice that we are also specifying `as_node=\"ask_human\"`\n",
    "# This will apply this update as this node,\n",
    "# which will make it so that afterwards it continues as normal\n",
    "app.update_state(config, {\"messages\": tool_message}, as_node=\"ask_human\")\n",
    "\n",
    "# We can check the state\n",
    "# We can see that the state currently has the `agent` node next\n",
    "# This is based on how we define our graph, \n",
    "# where after the `ask_human` node goes (which we just triggered)\n",
    "# there is an edge to the `agent` node\n",
    "app.get_state(config).next"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a30c9fb-2a40-45cc-87ba-406c11c9f0cf",
   "metadata": {},
   "source": [
    "We can now tell the agent to continue. We can just pass in `None` as the input to the graph, since no additional input is needed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a9f599b5-1a55-406b-a76b-f52b3ca06975",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "[{'text': \"Thank you for letting me know that you're in San Francisco. Now, I'll use the search function to look up the weather in San Francisco.\", 'type': 'text'}, {'id': 'toolu_01BKzfQfHBSUcgTbz7c4gs8w', 'input': {'query': 'current weather in San Francisco'}, 'name': 'search', 'type': 'tool_use'}]\n",
      "Tool Calls:\n",
      "  search (toolu_01BKzfQfHBSUcgTbz7c4gs8w)\n",
      " Call ID: toolu_01BKzfQfHBSUcgTbz7c4gs8w\n",
      "  Args:\n",
      "    query: current weather in San Francisco\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: search\n",
      "\n",
      "[\"I looked up: current weather in San Francisco. Result: It's sunny in San Francisco, but you better look out if you're a Gemini \\ud83d\\ude08.\"]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Based on the search results, I can provide you with information about the current weather in San Francisco:\n",
      "\n",
      "The weather in San Francisco is currently sunny. \n",
      "\n",
      "It's worth noting that the search result included a playful comment about Geminis, but that's not directly related to the weather information you requested. If you need more specific details about the temperature, humidity, or forecast, please let me know, and I'd be happy to search for more detailed weather information for San Francisco.\n"
     ]
    }
   ],
   "source": [
    "for event in app.stream(None, config, stream_mode=\"values\"):\n",
    "    event[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6f972d1-3d99-4fc1-8b33-92b71e74835d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
